#include "refine/smooth_b2.h"
#include "refine/window_pass.h"
#include <limits>
#include <iostream>
#include <set>
#include "imglib/imageclass.h"
#include <imglib/bitmap.h>
#include <imglib/bmpmorph.h>
#include <glib/g3/smatrix.h>

using std::clog;
using std::endl;
using std::flush;

namespace refine {

namespace smooth_b2 {

const params default_params = {2, 2, 1, 100, 0.0f, 3.0f, false};

typedef grey_mask::mask::elem_type grey_mask_value;
const grey_mask_value grey_mask_max_value = std::numeric_limits<grey_mask_value>::max();	//255
const grey_mask_value grey_mask_min_value = std::numeric_limits<grey_mask_value>::min();	//0
const grey_mask_value grey_mask_avg_value = (grey_mask_max_value + grey_mask_min_value) / 2;//127

typedef unsigned uint;
typedef g3::basic_vector<uint>		vectoru;
typedef g3::basic_smatrix<uint>		smatrixu;
typedef g3::basic_smatrix<float>	smatrixf;
typedef g3::basic_vector<float>		vectorf;

inline vectoru v(byte_color const & c)
	{ return vectoru(c.b, c.g, c.r); }

const static int fore_idx = -1; //this value in the map represents strict foreground pixel with which no variable is associated
const static int back_idx = -2;

/// maps pixel of unknown region to variable number
class varmap
{
public:
	varmap(img::bitmap const & strict_fore, img::bitmap const & strict_back) 
		: map(strict_fore.dim())
	{
		num = 0;
		for (img::image<int>::iterator v = map.begin(); v != map.end(); ++v) {
			if (strict_fore[v]) {
				assert (!strict_back[v]);
				*v = fore_idx;
			}
			else if (strict_back[v])
				*v = back_idx;
			else
				*v = num++;
		}
	}
	/// the total number of variables
	uint count() const
		{ return num; }
	int operator[] (size_point const & p) const
		{ return map[p]; }
	bool variable(size_point const & p) const
		{ return map[p] >= 0; }
private:
	img::image<int> map;
	uint num;
};

/// linear relation of central pixel's alpha is found in the form:
///   alpha^(n+1) = mean(alpha^n) + scalar_product(D^n, rel),
/// where
///   D^n = n * sum(alpha^n * I) - sum(alpha^n) * sum(I)
typedef vectorf relation;
typedef std::vector<relation> relations;

class relation_constructor
{
public:
	relation_constructor(byte_image const & im_, 
		varmap const & vmap_, relations & rels_, float noisiness_)
		: im(im_), vmap(vmap_), rels(rels_), noisiness(noisiness_)
	{ 
		rels.resize(vmap.count());
		remove_all(); 
		constructions_count = 0; 
	}
	size_extent dim() const
		{ return im.dim(); }
	uint count() const
		{ return constructions_count; }
	bool selected(size_point const & p) const
		{ return vmap.variable(p); }
	void add(size_point const & p) 
		{ add_sample(v(im[p])); }
	void remove(size_point const & p)
		{ remove_sample(v(im[p])); }
	void remove_all() {
		n = 0;
		I = vectoru::nill;
		II.zero();
	}
	void process(size_point const & p, size_rect const &) {
		++constructions_count;
		process_core(rels[vmap[p]], v(im[p]));
	}
private:
	void add_sample(vectoru const & i) {
		++n;
		I += i;
		II += tsquare(i);
	}
	void remove_sample(vectoru const & i) {
		--n;
		I -= i;
		II -= tsquare(i);
	}
	void process_core(relation & rel, vectorf const & i) 
	{
		smatrixf invR = (float(n) * smatrixf(II) - tsquare(vectorf(I))).inverse0(sqr(noisiness * n));
		rel = invR * (i - vectorf(I) * (float(1) / n));
	}
private:
	byte_image const & im;
	varmap const & vmap;
	relations & rels;
	uint constructions_count;
private:
	uint n;
	vectoru I;
	smatrixu II; ///< sum of (I * I^T)
	float noisiness;
};

//template <typename Raster, typename Matte>
class smoother
{
public:
	smoother(byte_image const & src_, varmap const & vmap_, relations const & rels_,
		grey_mask::mask const & matte_, grey_mask::mask & next_matte_,
		img::bitmap & modifiable_, float sharpness_)
		: src(src_), vmap(vmap_), rels(rels_), matte(matte_), next_matte(next_matte_), modifiable(modifiable_),
		  sharpness(sharpness_)
		{ clear(); }
	void clear() {
		checks = achanges = 0;
		remove_all();
	}
	size_extent dim() const
		{ return src.dim(); }
	uint check_count() const
		{ return checks; }
	uint change_count() const
		{ return achanges; }
	bool selected(size_point const & p) const
		{ return modifiable[p]; }
	void add(size_point const & p) 
		{ add_sample(v(src[p]), matte[p]); }
	void remove(size_point const & p) 
		{ remove_sample(v(src[p]), matte[p]); }
	void remove_all() { 
		 n = A = 0; 
		 I = AI = vectoru::nill;
	}
	void process(size_point const & p, size_rect const & wnd) 
	{
		++checks;

		vectorf D = float(n) * vectorf(AI) - float(A) * vectorf(I);
		float a = float(A) / n + scalar_product(D, rels[vmap[p]])
			 + 0.5f; //0.5 for proper rounding

		if (sharpness > 0) {
			int sum_a = matte[p];
			int count = 1;
			if (p.x > 0) {
				sum_a += matte(p.x-1, p.y);
				++count;
			}
			if (p.x + 1 < dim().sx) {
				sum_a += matte(p.x+1, p.y);
				++count;
			}
			if (p.y > 0) {
				sum_a += matte(p.x, p.y-1);
				++count;
			}
			if (p.y + 1 < dim().sy) {
				sum_a += matte(p.x, p.y+1);
				++count;
			}
			a += (sharpness * 2 * sum_a) / count - sharpness * grey_mask_max_value;
		}

		grey_mask_value g;
		if (a <= grey_mask_min_value)
			g = grey_mask_min_value;
		else if (a >= grey_mask_max_value)
			g = grey_mask_max_value;
		else 
			g = grey_mask_value(a);

		/*bool not_greater_exists = false;
		bool not_less_exists = false;
		size_point i;
		for (i.y = wnd.y1; i.y < wnd.y2; ++i.y)
			for (i.x = wnd.x1; i.x < wnd.x2; ++i.x) {
				if (!not_greater_exists && matte[i] <= g)
					not_greater_exists = true;
				if (!not_less_exists && matte[i] >= g)
					not_less_exists = true;
				if (not_greater_exists && not_less_exists)
					break;
			}*/

		if (/*not_greater_exists && not_less_exists &&*/ matte[p] != g) {
			++achanges;
			next_matte[p] = g;
		}
		else
			modifiable[p] = false;
	}
private:
	void add_sample(vectoru const & i, uint a) {
		++n;
		A += a;
		I += i;
		AI += a * i;
	}
	void remove_sample(vectoru const & i, uint a) {
		--n;
		A -= a;
		I -= i;
		AI -= a * i;
	}
private:
	byte_image const & src;
	varmap const & vmap;
	relations const & rels;
	grey_mask::mask const & matte;
	grey_mask::mask & next_matte;
	img::bitmap & modifiable;
	float sharpness;
	uint checks, achanges;
private:
	uint n, A;
	vectoru I, AI;
};

void perform_core(
	const byte_image & src, 
	grey_mask::mask & matte, 
	const img::bitmap & strict_fore, 
	const img::bitmap & strict_back,
	img::bitmap & modifiable, ///< map of pixels that may be changed during the following pass
	const params & p)
{
	if (!p.silent)
		clog << "Volosok-B2 matting..." << endl;

	size_extent dim = src.dim();
	assert (dim == matte.dim());
	assert (dim == modifiable.dim());
	int_rect frame(0, 0, dim.sx, dim.sy);

	varmap vmap(strict_fore, strict_back);
	relations rels(vmap.count());
	window_pass(relation_constructor(src, vmap, rels, p.noisiness), p.whs);

	grey_mask::mask next_matte = matte;

	for (uint pass = 0; pass < p.passes_num; ++pass) {
		if (!p.silent)
			clog << "pass " << pass+1 << '/' << p.passes_num << flush;

		img::combine(modifiable, strict_fore, img::bit_and_not());
		img::combine(modifiable, strict_back, img::bit_and_not());

		smoother sm(src, vmap, rels, matte, next_matte, modifiable, p.sharpness);
		window_pass(sm, p.whs);
		if (!p.silent)
			clog << ": checked = " << sm.check_count() << ", alpha changes = " << sm.change_count() << endl;
		if (sm.change_count() == 0)
			break;

		matte = next_matte;

		img::dilate8(modifiable, p.whs);
	}
}


void perform(
	const byte_image & src, 
	const bit_mask::mask & m, 
	grey_mask::mask & matte, 
	const params & p)
{
	size_extent dim = src.dim();
	assert (dim == m.dim());

	matte.resize(dim);
	for (grey_mask::mask::iterator i = matte.begin(); i != matte.end(); ++i)
		*i = m[i] ? grey_mask_max_value : grey_mask_min_value;

	img::bitmap strict_fore, strict_back;
	img::erode4(strict_fore, m, p.fore_depth);
	img::dilate4(strict_back, m, p.back_depth);
	strict_back.inverse();

	img::bitmap modifiable = m; 
	{
		uint w = 1 + p.whs;
		img::dilate4(modifiable, w);
		img::bitmap tmp = m;
		tmp.inverse();
		img::dilate4(tmp, w);
		img::combine(modifiable, tmp, img::bit_and());
	}

	perform_core(src, matte, strict_fore, strict_back, modifiable, p);
}

void perform(
	const byte_image & src, 
	grey_mask::mask & matte, 
	const params & p)
{
	size_extent dim = matte.dim();

	img::bitmap not_min(dim), not_max(dim), strict_fore(dim), strict_back(dim);
	for (grey_mask::mask::const_iterator i = matte.begin(); i != matte.end(); ++i) {
		not_min[i] = *i != grey_mask_min_value;
		not_max[i] = *i != grey_mask_max_value;
		strict_fore[i] = *i > grey_mask_avg_value;
	}

	strict_back = strict_fore;
	img::erode4(strict_fore, p.fore_depth);
	img::dilate4(strict_back, p.back_depth);
	strict_back.inverse();

	img::bitmap modifiable;
	{
		uint w = 1 + p.whs;
		img::dilate4(modifiable, not_min, w);
		img::bitmap tmp;
		img::dilate4(tmp, not_max, w);
		img::combine(modifiable, tmp, img::bit_and());
	}

	perform_core(src, matte, strict_fore, strict_back, modifiable, p);
}

void perform(
	argb_byte_image & src,
	const params & p)
{
	//!!! temporary solution
	byte_image src2 = src; 
	grey_mask::mask matte = img::projection(src, &argb_byte_color::a);

	perform(src2, matte, p);

	img::projection(src, &argb_byte_color::a) = matte;
}

} //namespace smooth_b2

} //namespace refine
